import numpy as np
from . import LocalSearchOperator, register_operator


class FourOptOperator(LocalSearchOperator):
    """4-opt局部搜索算子"""

    def __init__(self):
        self.name = "four_opt"
        self.supported_problems = ["tsp"]

    def generate_neighborhood(self, solution, problem_instance):
        """生成4-opt邻域（简化版，实际中通常不生成完整邻域）"""
        neighborhood = []
        n = len(solution)

        # 4-opt的完整邻域非常大(O(n^4))，通常不直接生成
        # 这里只生成少数几种常见的4-opt移动

        # 随机选择四个断点
        if n > 8:
            i = np.random.randint(1, n - 7)
            j = np.random.randint(i + 2, n - 5)
            k = np.random.randint(j + 2, n - 3)
            l = np.random.randint(k + 2, n - 1)

            # 生成几种常见的4-opt移动
            moves = self.generate_four_opt_moves(solution, i, j, k, l)
            neighborhood.extend(moves)

        return neighborhood

    def generate_four_opt_moves(self, solution, i, j, k, l):
        """生成基于四个断点的各种4-opt移动"""
        moves = []
        n = len(solution)

        # 分段: A = [0:i], B = [i:j], C = [j:k], D = [k:l], E = [l:]
        A = solution[0:i]
        B = solution[i:j]
        C = solution[j:k]
        D = solution[k:l]
        E = solution[l:]

        # 4-opt移动类型1: A + D + C + B + E (双桥移动)
        move1 = np.concatenate([A, D, C, B, E])
        moves.append(move1)

        # 4-opt移动类型2: A + C + B + D + E
        move2 = np.concatenate([A, C, B, D, E])
        moves.append(move2)

        # 4-opt移动类型3: A + B + D + C + E
        move3 = np.concatenate([A, B, D, C, E])
        moves.append(move3)

        # 4-opt移动类型4: A + C + D + B + E
        move4 = np.concatenate([A, C, D, B, E])
        moves.append(move4)

        # 4-opt移动类型5: A + D + B + C + E
        move5 = np.concatenate([A, D, B, C, E])
        moves.append(move5)

        # 添加一些反转段的变体
        move6 = np.concatenate([A, D[::-1], C, B, E])
        moves.append(move6)

        move7 = np.concatenate([A, D, C[::-1], B, E])
        moves.append(move7)

        move8 = np.concatenate([A, D, C, B[::-1], E])
        moves.append(move8)

        return moves

    def efficient_four_opt(self, solution, problem_instance, intensity=5):
        """高效的4-opt实现，直接计算成本变化"""
        n = len(solution)
        current_solution = solution.copy()
        current_cost = problem_instance.evaluate(current_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False

            # 尝试多个随机断点组合
            for _ in range(min(50, n)):  # 限制尝试次数
                # 随机选择四个断点
                if n > 8:
                    i = np.random.randint(1, n - 7)
                    j = np.random.randint(i + 2, n - 5)
                    k = np.random.randint(j + 2, n - 3)
                    l = np.random.randint(k + 2, n - 1)

                    # 计算原始连接的成本
                    orig_cost = (
                            problem_instance.distance_matrix[solution[i - 1]][solution[i]] +
                            problem_instance.distance_matrix[solution[j - 1]][solution[j]] +
                            problem_instance.distance_matrix[solution[k - 1]][solution[k]] +
                            problem_instance.distance_matrix[solution[l - 1]][solution[l]]
                    )

                    # 尝试不同的4-opt移动并计算成本变化
                    moves = self.generate_four_opt_moves(solution, i, j, k, l)

                    for move in moves:
                        # 计算新连接的成本
                        new_cost = (
                                problem_instance.distance_matrix[move[i - 1]][move[i]] +
                                problem_instance.distance_matrix[move[j - 1]][move[j]] +
                                problem_instance.distance_matrix[move[k - 1]][move[k]] +
                                problem_instance.distance_matrix[move[l - 1]][move[l]]
                        )

                        # 计算总成本变化
                        cost_change = new_cost - orig_cost

                        # 如果成本降低，则接受这个移动
                        if cost_change < 0:
                            current_solution = move
                            current_cost += cost_change
                            improved = True
                            break  # 找到一个改进就跳出内层循环

                    if improved:
                        break  # 找到一个改进就跳出尝试循环

            iterations += 1

        improvement = problem_instance.evaluate(solution) - current_cost
        return current_solution, improvement

    def __call__(self, solution, problem_instance, intensity=3):
        """应用4-opt局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        # 使用高效实现
        return self.efficient_four_opt(solution, problem_instance, intensity)


# 注册算子
register_operator("four_opt", FourOptOperator())